import torch
from torch.utils.data._utils.collate import default_collate

def merge(x, y):
    if isinstance(x, torch.Tensor):
        return torch.cat((x, y), dim=0)
    elif isinstance(x, list) or isinstance(x, tuple):
        return x + y
    else:
        raise NotImplementedError("merge not implemented for type {}".format(type(x)))

def collate_fn(batch, id_len=None):
    id_batch = batch[:id_len]
    od_batch = batch[id_len:]

    # in_domain_batch = default_collate(in_domain_batch)
    # out_domain_batch = default_collate(out_domain_batch)
    # # merge element after first element
    # assert len(in_domain_batch) <= len(out_domain_batch), "wrong setup with two datasets"    # out_domain_batch contains the crop lenth
    # other_elements = [merge(x, y) for x, y in zip(in_domain_batch[1:], out_domain_batch[1:])]
    # crop_pos = out_domain_batch[-1]
    return default_collate(id_batch), default_collate(od_batch)

